# coding=utf-8

#当前目录
#current_directory = '/root/paddlejob/workspace/code/'
current_directory = '/home/aistudio/mock'

# 数据集文件目录
#datasets_prefix = '/root/paddlejob/workspace/train_data/datasets/'
datasets_prefix='/home/aistudio/data/'

# 输出文件目录. 任务完成后平台会自动把该目录所有文件压缩为tar.gz包，用户可以通过「下载输出」可以将输出信息下载到本地.
#output_dir = "/root/paddlejob/workspace/output"
output_dir = "/home/aistudio/mock/output"

import os

if __name__ == "__main__":
    
    print(os.getcwd())

    print('1.解压源代码')   
    os.system("unzip PaddleOCR.zip -d ./")

    print('2.解压修改过的源代码，复制到指定位置') 
    os.system("unzip updated.zip")
    os.system("cp ./updated/requirements.txt ./PaddleOCR/.")
    os.system("cp ./updated/chinese20000.txt ./PaddleOCR/.")
    os.system("cp ./updated/rec_chinese_lite_train_v2.0.yml ./PaddleOCR/configs/rec/ch_ppocr_v2.0/.")
    os.system("cp ./updated/rec_chinese_common_train_v2.0.yml ./PaddleOCR/configs/rec/ch_ppocr_v2.0/.")


    print('3.进入到PaddleOCR目录')   
    os.chdir("PaddleOCR") 

    print('4.安装python第三方库')    
    os.system("pip install -r requirements.txt -i https://mirror.baidu.com/pypi/simple")

    print('5.解压第一个数据集')    
    os.system("mkdir train_data")
    os.system("unzip {}data112149/rotated_20000.zip -d ./train_data/".format(datasets_prefix))
    os.system("cp {}data112149/det_rotated_v.txt -d ./train_data/".format(datasets_prefix))
    os.system("cp {}data112149/rec_rotated_v.txt -d ./train_data/".format(datasets_prefix))

    print('6.解压第二个数据集')    
    os.system("mkdir test_data")
    os.system("unzip {}data112322/data.zip -d ./test_data/".format(datasets_prefix))
    os.system("cp {}data112322/det_rotated_v.txt -d ./test_data/".format(datasets_prefix))
    os.system("cp {}data112322/rec_rotated_v.txt -d ./test_data/".format(datasets_prefix))

    print('7.下载预训练集，解压')    
    os.system("mkdir pretrain_models")
    os.system("wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_rec_pre.tar")
    os.system("tar -xf ch_ppocr_mobile_v2.0_rec_pre.tar -C ./pretrain_models/")
    os.system("wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_server_v2.0_rec_pre.tar")
    os.system("tar -xf ch_ppocr_server_v2.0_rec_pre.tar -C ./pretrain_models/")


    print('8.开始训练')
    # 单卡
    # os.system("export CUDA_VISIBLE_DEVICES=0")
    # 四卡
    #os.system("export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7")
    

    os.system("python -m paddle.distributed.launch tools/train.py -c ./configs/rec/ch_ppocr_v2.0/rec_chinese_lite_train_v2.0.yml")
    os.system("mv ./output {}".format(output_dir))
